{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Tce3stUlHN0L"
   },
   "source": [
    "##### Copyright 2018 The TensorFlow Authors.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "tuOe1ymfHZPu"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MfBg1C5NB3X0"
   },
   "source": [
    "# Distributed training with TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xHxb-dlhMIzW"
   },
   "source": [
    "## Overview\n",
    "\n",
    "`tf.distribute.Strategy` is a TensorFlow API to distribute training \n",
    "across multiple GPUs, multiple machines or TPUs. Using this API, you can distribute your existing models and training code with minimal code changes.\n",
    "\n",
    "`tf.distribute.Strategy` has been designed with these key goals in mind:\n",
    "\n",
    "* Easy to use and support multiple user segments, including researchers, ML engineers, etc.\n",
    "* Provide good performance out of the box.\n",
    "* Easy switching between strategies.\n",
    "\n",
    "`tf.distribute.Strategy` can be used with a high-level API like [Keras](https://www.tensorflow.org/guide/keras), and can also be used to distribute custom training loops (and, in general, any computation using TensorFlow).\n",
    "\n",
    "In TensorFlow 2.x, you can execute your programs eagerly, or in a graph using [`tf.function`](function.ipynb). `tf.distribute.Strategy` intends to support both these modes of execution, but works best with `tf.function`. Eager mode is only recommended for debugging purpose and not supported for `TPUStrategy`. Although we discuss training most of the time in this guide, this API can also be used for distributing evaluation and prediction on different platforms.\n",
    "\n",
    "You can use `tf.distribute.Strategy` with very few changes to your code, because we have changed the underlying components of TensorFlow to become strategy-aware. This includes variables, layers, models, optimizers, metrics, summaries, and checkpoints.\n",
    "\n",
    "In this guide, we explain various types of strategies and how you can use them in different situations.\n",
    "\n",
    "Note: For a deeper understanding of the concepts, please watch [this deep-dive presentation](https://youtu.be/jKV53r9-H14). This is especially recommended if you plan to write your own training loop.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EVOZFbNgXghB"
   },
   "outputs": [],
   "source": [
    "# Import TensorFlow\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eQ1QESxxEbCh"
   },
   "source": [
    "## Types of strategies\n",
    "`tf.distribute.Strategy` intends to cover a number of use cases along different axes. Some of these combinations are currently supported and others will be added in the future. Some of these axes are:\n",
    "\n",
    "* *Synchronous vs asynchronous training:* These are two common ways of distributing training with data parallelism. In sync training, all workers train over different slices of input data in sync, and aggregating gradients at each step. In async training, all workers are independently training over the input data and updating variables asynchronously. Typically sync training is supported via all-reduce and async through parameter server architecture.\n",
    "* *Hardware platform:* You may want to scale your training onto multiple GPUs on one machine, or multiple machines in a network (with 0 or more GPUs each), or on Cloud TPUs.\n",
    "\n",
    "In order to support these use cases, there are six strategies available. In the next section we explain which of these are supported in which scenarios in TF 2.2 at this time. Here is a quick overview:\n",
    "\n",
    "| Training API          \t| MirroredStrategy  \t| TPUStrategy         \t| MultiWorkerMirroredStrategy     \t| CentralStorageStrategy          \t| ParameterServerStrategy  \t|\n",
    "|:-----------------------\t|:-------------------\t|:---------------------\t|:---------------------------------\t|:---------------------------------\t|:--------------------------\t|\n",
    "| **Keras API**             \t| Supported         \t|  Supported \t| Experimental support \t| Experimental support \t| Supported planned post 2.3 \t|\n",
    "| **Custom training loop**  \t| Supported \t| Supported   \t| Experimental support            \t| Experimental support           \t| Supported planned post 2.3         \t|\n",
    "| **Estimator API**         \t| Limited Support         \t| Not supported           \t| Limited Support                       \t| Limited Support                       \t| Limited Support                \t|\n",
    "\n",
    "Note: [Experimental support](https://www.tensorflow.org/guide/versions#what_is_not_covered) means the APIs are not covered by any compatibilities guarantees.\n",
    "\n",
    "Note: Estimator support is limited. Basic training and evaluation are experimental, and advanced features—such as scaffold—are not implemented. We recommend using Keras or custom training loops if a use case is not covered."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DoQKKK8dtfg6"
   },
   "source": [
    "### MirroredStrategy\n",
    "`tf.distribute.MirroredStrategy` supports synchronous distributed training on multiple GPUs on one machine. It creates one replica per GPU device. Each variable in the model is mirrored across all the replicas. Together, these variables form a single conceptual variable called `MirroredVariable`. These variables are kept in sync with each other by applying identical updates.\n",
    "\n",
    "Efficient all-reduce algorithms are used to communicate the variable updates across the devices.\n",
    "All-reduce aggregates tensors across all the devices by adding them up, and makes them available on each device.\n",
    "It’s a fused algorithm that is very efficient and can reduce the overhead of synchronization significantly. There are many all-reduce algorithms and implementations available, depending on the type of communication available between devices. By default, it uses NVIDIA NCCL as the all-reduce implementation. You can choose from a few other options we provide, or write your own.\n",
    "\n",
    "Here is the simplest way of creating `MirroredStrategy`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9Z4FMAY9ADxK"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wldY4aFCAH4r"
   },
   "source": [
    "This will create a `MirroredStrategy` instance which will use all the GPUs that are visible to TensorFlow, and use NCCL as the cross device communication.\n",
    "\n",
    "If you wish to use only some of the GPUs on your machine, you can do so like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nbGleskCACv_"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:1,/job:localhost/replica:0/task:0/device:GPU:0\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy(devices=[\"/gpu:0\", \"/gpu:1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8-KDnrJLAhav"
   },
   "source": [
    "If you wish to override the cross device communication, you can do so using the `cross_device_ops` argument by supplying an instance of `tf.distribute.CrossDeviceOps`. Currently,  `tf.distribute.HierarchicalCopyAllReduce` and `tf.distribute.ReductionToOneDevice` are two options other than `tf.distribute.NcclAllReduce` which is the default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6-xIOIpgBItn"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy(\n",
    "    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8Xc3gyo0Bejd"
   },
   "source": [
    "### MultiWorkerMirroredStrategy\n",
    "\n",
    "`tf.distribute.experimental.MultiWorkerMirroredStrategy` is very similar to `MirroredStrategy`. It implements synchronous distributed training across multiple workers, each with potentially multiple GPUs. Similar to `MirroredStrategy`, it creates copies of all variables in the model on each device across all workers.\n",
    "\n",
    "It uses [CollectiveOps](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/collective_ops.py) as the multi-worker all-reduce communication method used to keep variables in sync. A collective op is a single op in the TensorFlow graph which can automatically choose an all-reduce algorithm in the TensorFlow runtime according to hardware, network topology and tensor sizes.\n",
    "\n",
    "It also implements additional performance optimizations. For example, it includes a static optimization that converts multiple all-reductions on small tensors into fewer all-reductions on larger tensors. In addition, we are designing it to have a plugin architecture - so that in the future, you will be able to plugin algorithms that are better tuned for your hardware. Note that collective ops also implement other collective operations such as broadcast and all-gather.\n",
    "\n",
    "Here is the simplest way of creating `MultiWorkerMirroredStrategy`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "m3a_6ebbEjre"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)\n",
      "INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CollectiveCommunication.AUTO\n"
     ]
    }
   ],
   "source": [
    "multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bt94JBvhEr4s"
   },
   "source": [
    "`MultiWorkerMirroredStrategy` currently allows you to choose between two different implementations of collective ops.  `CollectiveCommunication.RING` implements ring-based collectives using gRPC as the communication layer.  `CollectiveCommunication.NCCL` uses [Nvidia's NCCL](https://developer.nvidia.com/nccl) to implement collectives.  `CollectiveCommunication.AUTO` defers the choice to the runtime.  The best choice of collective implementation depends upon the number and kind of GPUs, and the network interconnect in the cluster. You can specify them in the following way:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QGX_QAEtFQSv"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)\n",
      "INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CollectiveCommunication.NCCL\n"
     ]
    }
   ],
   "source": [
    "multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(\n",
    "    tf.distribute.experimental.CollectiveCommunication.NCCL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0JiImlw3F77E"
   },
   "source": [
    "One of the key differences to get multi worker training going, as compared to multi-GPU training, is the multi-worker setup. The `TF_CONFIG` environment variable is the standard way in TensorFlow to specify the cluster configuration to each worker that is part of the cluster. Learn more about [setting up TF_CONFIG](#TF_CONFIG)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "E20tG21LFfv1"
   },
   "source": [
    "Note: This strategy is [`experimental`](https://www.tensorflow.org/guide/versions#what_is_not_covered) as we are currently improving it and making it work for more scenarios. As part of this, please expect the APIs to change in the future."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "45H0Wa8WKI8z",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "### CentralStorageStrategy\n",
    "`tf.distribute.experimental.CentralStorageStrategy` does synchronous training as well. Variables are not mirrored, instead they are placed on the CPU and operations are replicated across all local GPUs. If there is only one GPU, all variables and operations will be placed on that GPU.\n",
    "\n",
    "Create an instance of `CentralStorageStrategy` by:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "id": "rtjZOyaoMWrP",
    "jupyter": {
     "outputs_hidden": true,
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:GPU:0'\n"
     ]
    }
   ],
   "source": [
    "central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KY1nJHNkMl7b",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "This will create a `CentralStorageStrategy` instance which will use all visible GPUs and CPU. Update to variables on replicas will be aggregated before being applied to variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aAFycYUiNCUb",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "Note: This strategy is [`experimental`](https://www.tensorflow.org/guide/versions#what_is_not_covered) as we are currently improving it and making it work for more scenarios. As part of this, please expect the APIs to change in the future."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3ZLBhaP9NUNr",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "### ParameterServerStrategy\n",
    "`tf.distribute.experimental.ParameterServerStrategy` supports parameter servers training on multiple machines. In this setup, some machines are designated as workers and some as parameter servers. Each variable of the model is placed on one parameter server. Computation is replicated across all GPUs of all the workers.\n",
    "\n",
    "In terms of code, it looks similar to other strategies:\n",
    "```\n",
    "ps_strategy = tf.distribute.experimental.ParameterServerStrategy()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zr1wPHYvOH0N",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "For multi worker training, `TF_CONFIG` needs to specify the configuration of parameter servers and workers in your cluster, which you can read more about in [TF_CONFIG below](#TF_CONFIG).\n",
    "\n",
    "Note: This strategy only works with the Estimator API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t2XUdmIxKljq",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "### Other strategies\n",
    "\n",
    "In addition to the above strategies, there are two other strategies which might be useful for prototyping and debugging when using `tf.distribute` APIs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UD5I1beTpc7a"
   },
   "source": [
    "#### Default Strategy\n",
    "\n",
    "Default strategy is a distribution strategy which is present when no explicit distribution strategy is in scope. It implements the `tf.distribute.Strategy` interface but is a pass-through and provides no actual distribution. For instance, `strategy.run(fn)` will simply call `fn`. Code written using this strategy should behave exactly as code written without any strategy. You can think of it as a \"no-op\" strategy.\n",
    "\n",
    "Default strategy is a singleton - and one cannot create more instances of it. It can be obtained using `tf.distribute.get_strategy()` outside any explicit strategy's scope (the same API that can be used to get the current strategy inside an explicit strategy's scope). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ibHleFOOmPn9"
   },
   "outputs": [],
   "source": [
    "default_strategy = tf.distribute.get_strategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EkxPl_5ImLzc"
   },
   "source": [
    "This strategy serves two main purposes:\n",
    "\n",
    "* It allows writing distribution aware library code unconditionally. For example, in optimizer, we can do `tf.distribute.get_strategy()` and use that strategy for reducing gradients - it will always return a strategy object on which we can call the reduce API.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WECeRzUdT6bU"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In optimizer or other library code\n",
    "# Get currently active strategy\n",
    "strategy = tf.distribute.get_strategy()\n",
    "strategy.reduce(\"SUM\", 1., axis=None)  # reduce some values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JURbH-pUT51B"
   },
   "source": [
    "* Similar to library code, it can be used to write end users' programs to work with and without distribution strategy, without requiring conditional logic. A sample code snippet illustrating this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "O4Vmae5jmSE6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>\n"
     ]
    }
   ],
   "source": [
    "if tf.config.list_physical_devices('gpu'):\n",
    "  strategy = tf.distribute.MirroredStrategy()\n",
    "else:  # use default strategy\n",
    "  strategy = tf.distribute.get_strategy() \n",
    "\n",
    "with strategy.scope():\n",
    "  # do something interesting\n",
    "  print(tf.Variable(1.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kTzsqN4lmJ0d",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "#### OneDeviceStrategy\n",
    "`tf.distribute.OneDeviceStrategy` is a strategy to place all variables and computation on a single specified device. \n",
    "\n",
    "```\n",
    "strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
    "```\n",
    "\n",
    "This strategy is distinct from the default strategy in a number of ways. In default strategy, the variable placement logic remains unchanged when compared to running TensorFlow without any distribution strategy. But when using `OneDeviceStrategy`, all variables created in its scope are explicitly placed on the specified device. Moreover, any functions called via `OneDeviceStrategy.run` will also be placed on the specified device. \n",
    "\n",
    "Input distributed through this strategy will be prefetched to the specified device. In default strategy, there is no input distribution.\n",
    "\n",
    "Similar to the default strategy, this strategy could also be used to test your code before switching to other strategies which actually distribute to multiple devices/machines. This will exercise the distribution strategy machinery somewhat more than default strategy, but not to the full extent as using `MirroredStrategy` or `TPUStrategy` etc. If you want code that behaves as if no strategy, then use default strategy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hQv1lm9UPDFy",
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "So far we've talked about what are the different strategies available and how you can instantiate them. In the next few sections, we will talk about the different ways in which you can use them to distribute your training. We will show short code snippets in this guide and link off to full tutorials which you can run end to end."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_mcuy3UhPcen"
   },
   "source": [
    "## Using `tf.distribute.Strategy` with `tf.keras.Model.fit`\n",
    "We've integrated `tf.distribute.Strategy` into `tf.keras` which is TensorFlow's implementation of the\n",
    "[Keras API specification](https://keras.io). `tf.keras`  is a high-level API to build and train models. By integrating into `tf.keras` backend, we've made it seamless for you to distribute your training written in the Keras training framework using `model.fit`.\n",
    "\n",
    "Here's what you need to change in your code:\n",
    "\n",
    "1. Create an instance of the appropriate `tf.distribute.Strategy`.\n",
    "2. Move the creation of Keras model, optimizer and metrics inside `strategy.scope`.\n",
    "\n",
    "We support all types of Keras models - sequential, functional and subclassed.\n",
    "\n",
    "Here is a snippet of code to do this for a very simple Keras model with one dense layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gbbcpzRnPZ6V"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "\n",
    "with mirrored_strategy.scope():\n",
    "  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])\n",
    "\n",
    "model.compile(loss='mse', optimizer='sgd')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "773EOxCRVlTg"
   },
   "source": [
    "In this example we used `MirroredStrategy` so we can run this on a machine with multiple GPUs. `strategy.scope()` indicates to Keras which strategy to use to distribute the training. Creating models/optimizers/metrics inside this scope allows us to create distributed variables instead of regular variables. Once this is set up, you can fit your model like you would normally. `MirroredStrategy` takes care of replicating the model's training on the available GPUs, aggregating gradients, and more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZMmxEFRTEjH5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "10/10 [==============================] - 0s 1ms/step - loss: 0.4668\n",
      "Epoch 2/2\n",
      "10/10 [==============================] - 0s 1ms/step - loss: 0.2063\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "10/10 [==============================] - 0s 1ms/step - loss: 0.1281\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.1281416267156601"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)\n",
    "model.fit(dataset, epochs=2)\n",
    "model.evaluate(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nofTLwyXWHK8"
   },
   "source": [
    "Here we used a `tf.data.Dataset` to provide the training and eval input. You can also use numpy arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Lqgd9SdxW5OW"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "10/10 [==============================] - 0s 1ms/step - loss: 0.0912\n",
      "Epoch 2/2\n",
      "10/10 [==============================] - 0s 1ms/step - loss: 0.0403\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7fa07027b2e8>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "inputs, targets = np.ones((100, 1)), np.ones((100, 1))\n",
    "model.fit(inputs, targets, epochs=2, batch_size=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "Tce3stUlHN0L"
   ],
   "name": "distributed_training.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
